import json
import sys
import re

def adjust_list(lst):
    if len(lst) == 0:
        return ['x', 'x', 'x', 'x', 'x']
    return lst

def read_jsonl2txt(file_path):
    data = []
    with open(file_path, 'r', encoding='utf-8') as f, open(file_path.replace('.jsonl', '.txt'), 'w', encoding='utf-8') as txt_file:
        for line in f:
            try:
                obj = json.loads(line)
                data.append(obj)
                prediction = obj['predict'].replace('\n', '')
                pattern = r"<answer>(.*?)</answer>"
                grpo_match = re.findall(pattern, prediction, re.DOTALL)
                if grpo_match:
                    kps = grpo_match[0].strip().split(',')
                    prediction = [kp.strip() for kp in kps if len(kp.strip())]
                elif ';' in prediction:
                    prediction = prediction.split(';')
                    prediction = [pred.strip() for pred in prediction if len(pred.strip())]
                else:
                    prediction = prediction.split(',')
                    prediction = [pred.strip() for pred in prediction if len(pred.strip())]
                predicts = adjust_list(prediction)
                pred_text = ';'.join(predicts)
                txt_file.write(pred_text + '\n')
            except json.JSONDecodeError as e:
                print(f"parse error: {e} at line: {line}")
    return data

pred_file = sys.argv[1]
data = read_jsonl2txt(pred_file)